package site.wanjiahao.trie;

import java.util.HashMap;
import java.util.Map;

public class PrefixTree<V> implements Trie<V>{

    private int size;

    // 根节点
    private Node<V> root;

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public void clear() {
        size = 0;
        root = null;
    }

    // 存放当前单词对应的次数
    @Override
    public V add(String key, V value) {
        // 非空检查
        checkPrefixIsLegal(key);
        if (root == null) {
            root = new Node<>(null, null);
        }
        Node<V> node = root;
        int len = key.length();
        int end = len - 1;

        for (int i = 0; i < len; i++) {
            char cur = key.charAt(i);
            // 新添加根对应的子节点
            Map<Character, Node<V>> children =  node.children;
            if (children == null) {
                children = new HashMap<>();
                node.children = children;
            };
            // 添加
            Node<V> parent = node;
            node = children.get(cur);
            if (node == null) {
                node = new Node<>(parent, cur);
                children.put(cur, node);
            }
            if (i == end) {
                node.word = true;
                node.value = value;
            }
        }
        size++;
        return node.value;
    }

    @Override
    public V get(String key) {
        Node<V> node = node(key);
        return (node != null && node.word)? node.value: null;
    }

    @Override
    public boolean contains(String key) {
        checkPrefixIsLegal(key);
        Node<V> node = node(key);
        return node != null && node.word;
    }

    @Override
    public V remove(String key) {
        checkPrefixIsLegal(key);
        Node<V> node = node(key);
        if (node == null) return null;
        V oldValue = node.value;
        if (node.children != null && !node.children.isEmpty()) {
            node.word = false;
            size--;
            return oldValue;
        }
        do {
            Node<V> parent = node.parent;
            // 删除
            parent.children.remove(node.character);
            // 判断parent是否还有子节点
            if (parent.word) {
                break;
            }
            if (parent.children != null && !parent.children.isEmpty()) {
                break;
            }
            node = parent;
        } while (node != null);
        size--;
        return oldValue;
    }

    @Override
    public boolean startsWith(String prefix) {
        checkPrefixIsLegal(prefix);
        return node(prefix) != null;
    }

    // 根据前缀，查询对应的节点
    private Node<V> node(String prefix) {
        if (root == null) return null;
        Node<V> node = root;
        for (int i = 0; i < prefix.length(); i++) {
            char cur = prefix.charAt(i);
            Map<Character, Node<V>> children = node.children;
            if (children == null || children.size() == 0) return null;
            node = children.get(cur);
            if (node == null) return null;
        }
        return node;
    }

    private void checkPrefixIsLegal(String prefix) {
        if (prefix == null || prefix.length() == 0) throw new IllegalArgumentException("IllegalArgument !");
    }

    // 节点对象
    private static class Node<V> {

        // 当前单词的值
        private V value;

        // 当前的父节点
        private final Node<V> parent;

        // 当前节点下的孩子
        private Map<Character, Node<V>> children;

        // 是否为单词结尾
        private boolean word;

        // 当前字符
        private Character character;

        public Node(Node<V> parent, Character character) {
            this.character = character;
            this.parent = parent;
        }
    }
}
